import gymnasium as gym
import mani_skill2.envs
import os
import numpy as np
# from stable_baselines3 import PPO

os.environ["CUDA_VISIBLE_DEVICES"] = '0'

# env = gym.make("SpacePipeline-v1", obs_mode="rgbd", control_mode="base_pd_joint_vel_arm_pd_ee_delta_pose", render_mode="human")
# env = gym.make("SpaceCatchObj-v1", obs_mode="rgbd", control_mode="base_pd_joint_vel_arm_pd_ee_delta_pose", render_mode="human")
# env = gym.make("SpaceTaskTwo-v1", obs_mode="state", control_mode="base_pd_joint_vel_arm_pd_ee_delta_pose", render_mode="human")
# env = gym.make("SpaceMoveBox-v1", obs_mode="rgbd", control_mode="base_pd_joint_vel_arm_pd_ee_delta_pose", render_mode="human")

# env = gym.make("SpaceBenchmark-v0", obs_mode="state", control_mode="base_pd_joint_vel_arm_pd_ee_delta_pose", render_mode="human")
# env = gym.make("SpaceOpenCabinet-v0", obs_mode="rgbd", control_mode="pd_joint_delta_pos", render_mode="human")
# env = gym.make("SpaceGraspCube-v0", obs_mode="rgbd", control_mode="pd_joint_delta_pos", render_mode="human")
# env = gym.make("SpaceSyringe-v0", obs_mode="rgbd", control_mode="pd_joint_delta_pos", render_mode="human")
# env = gym.make("Pour-v0", obs_mode="rgbd", control_mode="pd_joint_delta_pos", render_mode="human")
env = gym.make("CocubeGripper-v0", obs_mode="rgbd", control_mode="pd_joint_vel", render_mode="human")


print("Observation space", env.observation_space)
print("Action space", env.action_space)

obs, _ = env.reset(seed=0) # reset with a seed for randomness

for i in range(1000000):
    env.render()

# terminated, truncated = False, False
# while not terminated and not truncated:
#     action = env.action_space.sample()
#     action_dim = env.agent.action_space.shape[0]
#     obs, reward, terminated, truncated, info = env.step(action)
#     env.render()  # a display is required to render
# env.close()


# action_dim = env.agent.action_space.shape[0]
# action_dim = env.agent.action_space.shape[0]

# obs, reward, terminated, truncated, info = env.step(np.zeros(action_dim))
# env.render()  # a display is required to render
# env.unwrapped.demo()
# while not env.unwrapped.viewer.closed:
#     env.unwrapped._scene.step()
#     env.unwrapped._scene.update_render()
#     env.unwrapped.viewer.render()
# obs, reward, terminated, truncated, info = env.step(np.zeros(action_dim))
# env.render()  # a display is required to render
# env.unwrapped.demo()
# while not env.unwrapped.viewer.closed:
#     env.unwrapped._scene.step()
#     env.unwrapped._scene.update_render()
#     env.unwrapped.viewer.render()